import os
import shutil
import random
from data_split_kfold import get_Murmur_locations
from patient_information import find_patient_files,load_patient_data,get_grade,get_murmur
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from python_speech_features import logfbank,fbank
from sklearn.model_selection import StratifiedKFold
from tqdm import tqdm
import wave
import librosa
import soundfile
import librosa.display
from segmentation import tsv_load,get_beat_time
from scipy import signal
import warnings
warnings.filterwarnings("ignore")

def FrameTimeC(frameNum, frameLen, inc, fs):
    ll = np.array([i for i in range(frameNum)])
    return ((ll - 1) * inc + frameLen / 2) / fs


def FrequencyScale(nfilt, fs):
    high_freq_mel = (2595 * np.log10(1 + (fs / 2) / 700))  # 求最高hz频率对应的mel频率
    # 我们要做40个滤波器组，为此需要42个点，这意味着在们需要low_freq_mel和high_freq_mel之间线性间隔40个点
    mel_points = np.linspace(0, high_freq_mel, nfilt + 2)  # 在mel频率上均分成42个点
    hz_points = (700 * (10 ** (mel_points / 2595) - 1))  # 将mel频率再转到hz频率
    return hz_points


# 训练集数据增强后的特征（变速）
def time_change_feature(data_directory):
    logmelspec = list()
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':

            the_location = root.split("_")[1].strip()
            the_id = root.split("_")[0].strip()
            # 将不存在杂音的听诊位置标为absent
            label = 'Absent'
            txt_data = load_patient_data(os.path.join(data_directory, the_id + '.txt'))
            murmur_locations = (get_Murmur_locations(txt_data)).split("+")
            for i in range(len(murmur_locations)):
                if the_location == murmur_locations[i]:
                    label = root.split("_")[2].strip()

            x, fs = librosa.load(os.path.join(data_directory, f), sr=4000)
            x = max_normal(x)
            # b, a = signal.butter(4, [0.01, 0.4], 'bandpass')  # 配置滤波器 8 表示滤波器的阶数
            # x = signal.filtfilt(b, a, x)

            if label == 'Absent':
                # xx = max_normal(x)
                fbank_feat = logfbank(x, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0, highfreq=800).T
                logmelspec.append(fbank_feat)

                # for i in range(2):
                #     fbank_feat, _ = fbank(x, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                #                           highfreq=800)
                #     fbank_feat = fbank_feat.T
                #     mask_feat = np.log(mask_time(x, fbank_feat)+1e-5)
                #     logmelspec.append(mask_feat)

            if label == 'Soft':
                y = librosa.effects.time_stretch(x, rate=1.1)
                yy = normal_len(y, len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                      highfreq=800).T
                logmelspec.append(fbank_feat)
                y = librosa.effects.time_stretch(x, rate=1.2)
                yy = normal_len(y,len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0, highfreq=800).T
                logmelspec.append(fbank_feat)
                y = librosa.effects.time_stretch(x, rate=0.9)
                yy = normal_len(y, len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                      highfreq=800).T
                logmelspec.append(fbank_feat)
                # y = librosa.effects.time_stretch(x, rate=0.8)
                # yy = normal_len(y, len(x))
                # fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0, highfreq=800).T
                # logmelspec.append(fbank_feat)
                # # xx = max_normal(x)
                fbank_feat = logfbank(x, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0, highfreq=800).T
                logmelspec.append(fbank_feat)
                for i in range(3):
                    fbank_feat, _ = fbank(x, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                          highfreq=800)
                    fbank_feat = fbank_feat.T
                    mask_feat = np.log(mask_time(x, fbank_feat)+1e-5)
                    logmelspec.append(mask_feat)
            if label == 'Loud':  # 五倍
                y = librosa.effects.time_stretch(x, rate=1.1)
                yy = normal_len(y, len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                      highfreq=800).T
                logmelspec.append(fbank_feat)
                # y = librosa.effects.time_stretch(x, rate=1.15)
                # yy = normal_len(y, len(x))
                # fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                #                       highfreq=800).T
                # logmelspec.append(fbank_feat)
                y = librosa.effects.time_stretch(x, rate=1.2)
                yy = normal_len(y, len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                      highfreq=800).T
                logmelspec.append(fbank_feat)
                # y = librosa.effects.time_stretch(x, rate=1.25)
                # yy = normal_len(y, len(x))
                # fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                #                       highfreq=800).T
                # logmelspec.append(fbank_feat)
                # y = librosa.effects.time_stretch(x, rate=1.3)
                # yy = normal_len(y, len(x))
                # fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                #                       highfreq=800).T
                # logmelspec.append(fbank_feat)
                # y = librosa.effects.time_stretch(x, rate=0.95)
                # yy = normal_len(y, len(x))
                # fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                #                       highfreq=800).T
                # logmelspec.append(fbank_feat)
                y = librosa.effects.time_stretch(x, rate=0.9)
                yy = normal_len(y, len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                      highfreq=800).T
                logmelspec.append(fbank_feat)
                # y = librosa.effects.time_stretch(x, rate=0.85)
                # yy = normal_len(y, len(x))
                # fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                #                       highfreq=800).T
                # logmelspec.append(fbank_feat)
                y = librosa.effects.time_stretch(x, rate=0.8)
                yy = normal_len(y, len(x))
                fbank_feat = logfbank(yy, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                      highfreq=800).T
                logmelspec.append(fbank_feat)

                fbank_feat = logfbank(x, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0, highfreq=800).T
                logmelspec.append(fbank_feat)

                for i in range(3):
                    fbank_feat, _ = fbank(x, fs, winlen=0.025, winstep=0.0125, nfilt=32, nfft=512, lowfreq=0,
                                          highfreq=800)
                    fbank_feat = fbank_feat.T
                    mask_feat = np.log(mask_time(x, fbank_feat)+1e-5)
                    logmelspec.append(mask_feat)
        else:
            continue
    return np.array(logmelspec)


def get_time_change_label(data_directory):
    label = list()
    for f in sorted(os.listdir(data_directory)):
        root, extension = os.path.splitext(f)
        if extension == '.wav':
            the_location = root.split("_")[1].strip()
            the_id = root.split("_")[0].strip()
            # 将不存在杂音的听诊位置标为absent
            the_label = 'Absent'
            txt_data = load_patient_data(os.path.join(data_directory, the_id + '.txt'))
            murmur_locations = (get_Murmur_locations(txt_data)).split("+")
            for i in range(len(murmur_locations)):
                if the_location == murmur_locations[i]:
                    the_label = root.split("_")[2].strip()

            if the_label == 'Absent':
                grade = 0
                label.append(grade)
            elif the_label == 'Soft':
                grade = 1
                for i in range(7):
                    label.append(grade)
            elif the_label == 'Loud':
                grade = 2
                for i in range(8):
                    label.append(grade)

    return np.array(label)


def max_normal(x):
    x = x - np.mean(x)
    x = x / np.max(np.abs(x))
    return x

def normal_len(x,n):
    if len(x) > n:
        x = x[0:n]
    else:
        b = np.zeros(n - len(x))
        x = np.append(x, b)
    return x
def mask_time(x,fbank_feat):
    mask_fbank = fbank_feat
    max_time = min(int((len(x) - 100) / 50)+1,len(mask_fbank[0]))
    #生成掩码的列数
    num_columns_to_zero = random.randint(0, int(max_time/4))
    columns_to_zero = np.random.choice(max_time, num_columns_to_zero, replace=False)
    # 创建一个布尔掩码，其中被选择的列为True，其余列为False
    column_mask = np.zeros(len(mask_fbank[0]), dtype=bool)
    column_mask[columns_to_zero] = True
    # 将被选择的列的数据设置为零
    mask_fbank[:, column_mask] = 0
    return mask_fbank

if __name__ == "__main__":
    data_directory = "F:/heart_data/2022_challenge_new/the-circor-digiscope-phonocardiogram-dataset-1.0.3/training_data"
    out_directory="data_5fold_new2"
    # random_state=5678
    stratified_features = ["Absent", "Soft", "Loud"]
    files = os.listdir(out_directory)
    for f in files:
        # if f != '说明.txt':
        train_data_directory = os.path.join(out_directory, f, "train_data")
        # vali_data_directory = os.path.join(out_directory, f, "vali_data")
        # test_data_directory = os.path.join(out_directory,  "test_data")
        label_directory = os.path.join(out_directory, f, "label")
        logmel_directory = os.path.join(out_directory, f, "logmel")

        train_feature = time_change_feature(train_data_directory)
        np.save(logmel_directory + r'/train_feature_expand.npy', train_feature)

        train_label = get_time_change_label(train_data_directory)
        np.save(label_directory + r'/train_label_expand.npy', train_label)